package cn.demoncat.util.sec;

import static java.lang.String.format;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Reader;
import java.io.StringReader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.KeyPair;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;

import cn.demoncat.util.sec.constant.SecurityConstant.KeyStoreType;

/**
 * KeyStore
 * 
 * @see <a href="https://github.com/heroku/env-keystore">env-keystore</a>
 * 
 * @author 延晓磊
 *
 * @since 2020年6月8日
 */
public class BaseKeyStore {

	// 密码
	private final String password;
	// 证书库
	private final KeyStore keystore;

	// 证书类型：JKS, JCEKS, PKCS12, BKS，UBER
	public static final String DEFAULT_TYPE = KeyStoreType.PKCS12;

	/**
	 * 创建密钥库KeyStore
	 * 
	 * @param key		密钥：.key
	 * @param cert		证书：.pem
	 * @param password	密码
	 * 
	 * @author 延晓磊
	 * @throws Exception 
	 *
	 * @since 2020年6月8日
	 */
	public BaseKeyStore(Reader key, Reader cert, String password) throws Exception {
		this.password = password;
		this.keystore = createKeyStore(key, cert, password);
	}
	
	/**
	 * 创建密钥库KeyStore
	 * 
	 * @param key		密钥：.key
	 * @param cert		证书：.pem
	 * @param password	密码
	 * 
	 * @author 延晓磊
	 * @throws Exception 
	 *
	 * @since 2020年6月8日
	 */
	public BaseKeyStore(String key, String cert, String password) throws Exception {
		this(new StringReader(key), new StringReader(cert), password);
	}

	/**
	 * 创建信任库TrustStore
	 * 
	 * @param cert		证书：.pem
	 * @param password	密码
	 * 
	 * @author 延晓磊
	 * @throws Exception 
	 *
	 * @since 2020年6月8日
	 */
	public BaseKeyStore(Reader cert, String password) throws Exception {
		this.password = password;
		this.keystore = createTrustStore(cert);
	}
	
	/**
	 * 创建信任库TrustStore
	 * 
	 * @param cert		证书：.pem
	 * @param password	密码
	 * 
	 * @author 延晓磊
	 * @throws Exception 
	 *
	 * @since 2020年6月8日
	 */
	public BaseKeyStore(String cert, String password) throws Exception {
		this(new StringReader(cert), password);
	}

	/**
	 * 创建信任库TrustStore
	 * 
	 * @param cert		证书：.pem
	 * @param password	密码
	 * @param issuers	证书链
	 * 
	 * @author 延晓磊
	 * @throws Exception 
	 *
	 * @since 2020年6月8日
	 */
	public BaseKeyStore(Reader cert, String password, X509Certificate[] issuers) throws Exception {
		this.password = password;
		this.keystore = createTrustStore(cert, issuers);
	}
	
	/**
	 * 创建信任库TrustStore
	 * 
	 * @param cert		证书：.pem
	 * @param password	密码
	 * @param issuers	证书链
	 * 
	 * @author 延晓磊
	 * @throws Exception 
	 *
	 * @since 2020年6月8日
	 */
	public BaseKeyStore(String cert, String password, X509Certificate[] issuers) throws Exception {
		this(new StringReader(cert), password, issuers);
	}

	/**
	 * 获取密码
	 * 
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public String getPassword() {
		return password;
	}

	/**
	 * 获取密钥库
	 * 
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public KeyStore getKeystore() {
		return keystore;
	}

	/**
	 * 获取证书类型
	 * 
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public static String getDefaultType() {
		return DEFAULT_TYPE;
	}

	/**
	 * 读取Store
	 * 
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public InputStream toInputStream() throws Exception {
		return new ByteArrayInputStream(toBytes());
	}

	/**
	 * 读取Store
	 * 
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public byte[] toBytes() throws Exception {
		ByteArrayOutputStream bos = new ByteArrayOutputStream();
		this.store(bos);
		bos.close();
		return bos.toByteArray();
	}

	/**
	 * 输出Store
	 * 
	 * @param out
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public void store(OutputStream out) throws Exception {
		this.keystore.store(out, password.toCharArray());
	}

	/**
	 * 保存Store
	 * 
	 * @param path
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public void store(Path path) throws Exception {
		Files.write(path, toBytes());
	}
	
	/**
	 * 保存Store
	 * 
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public File storeTemp() throws Exception {
		File temp = File.createTempFile("env-keystore", getDefaultType().toLowerCase());
		store(temp.toPath());
		return temp;
	}
	
	/**
	 * 保存Store
	 * 
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	public void asFile(Consumer<File> c) throws Exception {
		File temp = storeTemp();
		c.accept(temp);
		Files.delete(temp.toPath());
	}

	/**
	 * 创建密钥库
	 * 
	 * @param keyReader
	 * @param certReader
	 * @param password
	 * @return
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	protected static KeyStore createKeyStore(final Reader keyReader, final Reader certReader, final String password) throws Exception {
		PrivateKey key = getPrivateKeyFromPEM(keyReader);
		keyReader.close();

		PEMParser parser = new PEMParser(certReader);
		KeyStore ks = KeyStore.getInstance(DEFAULT_TYPE);
		ks.load(null);

		List<X509Certificate> certificates = new ArrayList<>();
		X509Certificate certificate;
		while ((certificate = parseCert(parser)) != null) {
			certificates.add(certificate);
		}

		ks.setKeyEntry("alias", key, password.toCharArray(), certificates.toArray(new X509Certificate[] {}));
		parser.close();
		return ks;
	}

	/**
	 * 获取私钥
	 * 
	 * @param keyReader	私钥文件
	 * @return
	 * @throws IOException
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	protected static PrivateKey getPrivateKeyFromPEM(final Reader keyReader) throws IOException {
		final JcaPEMKeyConverter jcaPEMKeyConverter = new JcaPEMKeyConverter();
		final PEMParser pem = new PEMParser(keyReader);
		PrivateKey key;
		Object pemContent = pem.readObject();
		if (pemContent instanceof PEMKeyPair) {
			PEMKeyPair pemKeyPair = (PEMKeyPair) pemContent;
			KeyPair keyPair = jcaPEMKeyConverter.getKeyPair(pemKeyPair);
			key = keyPair.getPrivate();
		} else if (pemContent instanceof PrivateKeyInfo) {
			PrivateKeyInfo privateKeyInfo = (PrivateKeyInfo) pemContent;
			key = jcaPEMKeyConverter.getPrivateKey(privateKeyInfo);
		} else {
			pem.close();
			throw new IllegalArgumentException("Unsupported private key format '" + pemContent.getClass().getSimpleName() + '"');
		}
		pem.close();
		return key;
	}

	/**
	 * 创建信任库
	 * 
	 * @param certReader
	 * @return
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	protected static KeyStore createTrustStore(final Reader certReader) throws Exception {
		return createTrustStore(certReader, new X509Certificate[] {});
	}

	/**
	 * 创建信任库
	 * 
	 * @param certReader
	 * @param acceptedIssuers
	 * @return
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	protected static KeyStore createTrustStore(final Reader certReader, final X509Certificate[] acceptedIssuers) throws Exception {
		PEMParser parser = new PEMParser(certReader);
		KeyStore ks = KeyStore.getInstance(DEFAULT_TYPE);
		ks.load(null);

		int i = 0;
		for (X509Certificate certificate : acceptedIssuers) {
			ks.setCertificateEntry(format("alias%d", i), certificate);
			i += 1;
		}

		X509Certificate certificate;
		while ((certificate = parseCert(parser)) != null) {
			ks.setCertificateEntry(format("alias%d", i), certificate);
			i += 1;
		}
		parser.close();
		return ks;
	}

	/**
	 * 解析证书
	 * 
	 * @param parser
	 * @return
	 * @throws Exception
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年6月8日
	 */
	protected static X509Certificate parseCert(PEMParser parser) throws Exception {
		X509CertificateHolder certHolder = (X509CertificateHolder) parser.readObject();
		if (certHolder == null) {
			return null;
		}
		return new JcaX509CertificateConverter().getCertificate(certHolder);
	}
}
